import torch
import torch.nn as nn
from transformers import AutoModel, AutoConfig
from transformers.models.deberta_v2.configuration_deberta_v2 import DebertaV2Config

"""均值池化，对变长序列的隐藏状态进行掩码加权平均，生成固定长度的句向量"""
class MeanPooling(nn.Module):
    def __init__(self):
        super(MeanPooling, self).__init__() # 继承nn.Moudle成为pytorch标准神经网络模块

    """
    last_hidden_state：形状为 (batch_size, seq_len, hidden_size) 的三维张量，表示Transformer最后一层的隐藏状态。
    attention_mask：形状为 (batch_size, seq_len) 的二维张量，数据类型通常为 torch.long 或 torch.bool，区分有效token和填充token。
    输入为Transformer模型的“最后一层隐藏状态和对应的注意力掩码”，输出为每个序列的均值池化结果，即每个样本的token相加取平均。
    """
    def forward(self, last_hidden_state, attention_mask):
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
        sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded, 1)
        sum_mask = input_mask_expanded.sum(1)
        sum_mask = torch.clamp(sum_mask, min=1e-9)
        mean_embeddings = sum_embeddings / sum_mask
        return mean_embeddings


class Model(nn.Module):
    def __init__(self, cfg, config_path=None, pretrained=False, is_training = True):
        super().__init__()
        self.cfg = cfg
        self.dropout = cfg.DROPOUT_RATE
        self.is_training = is_training

        torch.serialization.add_safe_globals([DebertaV2Config])

        if config_path is None:
            self.config = AutoConfig.from_pretrained(cfg.MODEL, output_hidden_states=True)
            self.config.hidden_dropout = 0.
            self.config.hidden_dropout_prob = 0.
            self.config.attention_dropout = 0.
            self.config.attention_probs_dropout_prob = 0.

        else:
            self.config = torch.load(config_path)

        if pretrained:
            self.model = AutoModel.from_pretrained(cfg.MODEL, config=self.config)
        else:
            self.model = AutoModel.from_config(self.config)


        if self.cfg.GRADIENT_CHECKPOINTING:
            self.model.gradient_checkpointing_enable()

        self.pool = MeanPooling() # 均值池化 [batch_size, hidden_size]

        self.head = nn.Sequential(
            nn.Linear(self.config.hidden_size, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(self.dropout),
            nn.Linear(512, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(self.dropout),
            nn.Linear(64, 1)
        )
        self._init_weights(self.head)

    def _init_weights(self, module):

        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def feature(self, inputs):

        outputs = self.model(**inputs)
        last_hidden_states = outputs[0] # 提取最后隐藏状态
        feature = self.pool(last_hidden_states, inputs['attention_mask'])
        return feature

    def forward(self, inputs):

        feature = self.feature(inputs) # 特征提取
        output = self.head(feature) # 回归预测
        return output


